package com.boarsoft.soagov.javaagent;

import java.io.File;
import java.io.FileFilter;
import java.io.FileInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.boarsoft.soagov.javaagent.ClassTransformer;
import com.sun.tools.attach.VirtualMachine;
import com.sun.tools.attach.VirtualMachineDescriptor;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.LoaderClassPath;

public class AgentMain implements ClassFileTransformer {
	private static final Logger log = LoggerFactory.getLogger(AgentMain.class);

	private static List<String> ignoreLoaders = new ArrayList<>();
	private static final Map<ClassLoader, ClassPool> classPoolMap = new ConcurrentHashMap<>();
	private static final List<ClassTransformer> transformers = new ArrayList<>();

	public AgentMain() {
		ignoreLoaders.add("sun.reflect.DelegatingClassLoader");
		ignoreLoaders.add("org.apache.catalina.loader.StandardClassLoader");
		ignoreLoaders.add("javax.management.remote.rmi.NoCallStackClassLoader");
		ignoreLoaders.add("com.alibaba.fastjson.util.ASMClassLoader");
		// TODO 从指定目录读取java文件，并装载
		ClassLoader cl = AgentMain.class.getClassLoader();
		ClassPool cp = new ClassPool();
		cp.insertClassPath(new LoaderClassPath(cl));
		classPoolMap.put(cl, cp);
		File td = new File(System.getProperty("transformer.dir", "transformer"));
		log.info("Scan transformers in {}", td.getAbsolutePath());
		if (td.exists() && td.isDirectory() && td.canRead()) {
			File[] fa = td.listFiles(new FileFilter() {
				@Override
				public boolean accept(File f) {
					return f.exists() && f.isFile() && f.canRead()//
							&& f.getName().endsWith(".class");
				}
			});
			for (File f : fa) {
				try {
					log.info("Read java class file {}", f.getAbsolutePath());
					CtClass cc = cp.makeClass(new FileInputStream(f));
					transformers.add((ClassTransformer) cc.toClass().newInstance());
				} catch (Exception e) {
					log.error("Can not make tranformer {}", f.getName(), e);
				}
			}
		}
	}

	@Override
	public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
			ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
		// 忽略特定的类
		if (className == null || loader == null) {
			return null;
		}
		if (className.indexOf('$') >= 0 || className.startsWith("java")//
				|| ignoreLoaders.contains(loader.getClass().getName())) {
			// log.info("Skip class {}", className);
			return null;
		}
		// 只转换符合条件的类
		// log.info("Try transform class {}", className);
		if (!classPoolMap.containsKey(loader)) {
			ClassPool classPool = new ClassPool();
			classPool.insertClassPath(new LoaderClassPath(loader));
			classPoolMap.put(loader, classPool);
		}
		ClassPool cp = classPoolMap.get(loader);
		for (ClassTransformer ct : transformers) {
			try {
				byte[] ba = ct.transform(cp, className);
				if (ba != null) {
					return ba;
				}
			} catch (Exception e) {
				log.error("Error on transform class {}", className, e);
				return null;
			}
		}
		return null;
	}

	public static void premain(String args, Instrumentation inst) {
		log.info("Premain args is {}", args);
		inst.addTransformer(new AgentMain(), true);
	}

	/**
	 * 
	 * @param args
	 * @param inst
	 */
	public static void agentmain(String args, Instrumentation inst) {
		log.info("Agentmain args is {}", args);
		inst.addTransformer(new AgentMain(), true);
	}

	/**
	 * 
	 * @param agentArgs
	 */
	public static void agentmain(String agentArgs) {
		// Nothing to do
	}

	public static void main(String[] args) throws Exception {
		// 获取当前系统中所有 运行中的 虚拟机
		log.info("Try attach JVM {} with {}", args[0], args[1]);
		List<VirtualMachineDescriptor> list = VirtualMachine.list();
		for (VirtualMachineDescriptor vmd : list) {
			// 如果虚拟机的名称为 xxx 则 该虚拟机为目标虚拟机，获取该虚拟机的 pid
			String vdn = vmd.displayName();
			// 然后加载 agent.jar 发送给该虚拟机
			log.info("Found JVM {}", vdn);
			if (vdn.endsWith(args[0])) {
				log.info("Attach JVM {}", vdn);
				VirtualMachine vm = VirtualMachine.attach(vmd.id());
				vm.loadAgent(args[1]);
				vm.detach();
				log.info("Detach JVM {}", vdn);
			}
		}
	}

}
